# -*- coding: utf-8 -*-
# !/usr/bin/env python
"""
-------------------------------------------------
   File Name：     train
   Description :   
   Author :       lth
   date：          2022/2/19
-------------------------------------------------
   Change Activity:
                   2022/2/19 6:13: create this script
-------------------------------------------------
"""
__author__ = 'lth'

import random
import warnings

import numpy as np
import torch
from PIL import ImageFont, ImageDraw, Image
from colorama import Fore
from torch import optim
from torch.backends import cudnn
from torch.utils.data import DataLoader
from tqdm import tqdm

from config import GetConfig
from datalist import CenterDataset, center_dataset_collate
from model import CenterNet
from utils import TargetEncode, CenterLoss, TargetDecode, get_image, weights_init


warnings.filterwarnings('ignore')
loss = 100000
font = ImageFont.truetype(font='model_data/simhei.ttf',
                          size=np.floor(3e-2 * 1024 + 0.5).astype('int32'))


def get_classes(class_path="./data/class.txt"):
    with open(class_path, 'r') as f:
        class_names = f.readlines()
    return [c.strip('\n') for c in class_names]


def line_process(lines):
    return_lines = []
    for line in lines:
        temp = line.split()[1:]
        ans = []
        for i in range(4, len(temp), 5):
            ans.append(temp[i])
        if len(set(ans)) == 1 and ans[0] == "14":
            continue
        else:
            return_lines.append(line)
    return return_lines


class CenterNetTrain(object):
    def __init__(self):
        self.args = GetConfig()

        print(f"-----------{self.args.project_name}-------------")

        # region 项目运行配置
        use_cuda = self.args.use_cuda and torch.cuda.is_available()
        self.device = torch.device('cuda' if use_cuda else 'cpu')
        kwargs = {'num_workers': 4, 'pin_memory': True} if use_cuda else {}
        # endregion

        # region 数据集dataloader
        self.annotations = self.args.annotations
        with open(self.annotations) as f:
            self.lines = f.readlines()

        self.lines = line_process(self.lines)

        random.shuffle(self.lines)

        self.num_val = int(len(self.lines) * self.args.val_num)
        self.num_train = len(self.lines) - self.num_val

        self.train_dataset = CenterDataset(self.lines, self.args.image_size)
        self.test_dataset = CenterDataset(self.lines, self.args.image_size)

        self.train_dataloader = DataLoader(self.train_dataset, batch_size=self.args.train_batch_size, shuffle=True,
                                           collate_fn=center_dataset_collate, **kwargs)
        self.test_dataloader = DataLoader(self.test_dataset, batch_size=self.args.test_batch_size, shuffle=True,
                                          collate_fn=center_dataset_collate, **kwargs)

        self.classes = get_classes(self.args.class_path)
        self.model = CenterNet(len(self.classes))

        if use_cuda:
            self.model = torch.nn.DataParallel(self.model, device_ids=range(torch.cuda.device_count())).to(self.device)
            cudnn.benchmark = True

        # ToDo : use resume or not
        if not self.args.resume:
            print("load the weight from pretrained-weight file")
            model_dict = self.model.state_dict()
            pretrained_dict = torch.load(self.args.pretrained_weight, map_location=self.device)['model_state_dict']
            pretrained_dict = {k: v for k, v in pretrained_dict.items() if np.shape(model_dict[k]) == np.shape(v)}
            model_dict.update(pretrained_dict)
            self.model.load_state_dict(model_dict)
            print("Finished to load the weight")
        else:
            print("train from scratch")
            # weights_init(self.model, init_type="normal", init_gain=0.01)

        self.optimizer = optim.AdamW(self.model.parameters(), self.args.lr, weight_decay=1e-4)
        self.scheduler = optim.lr_scheduler.StepLR(self.optimizer, 3, gamma=0.85)

        self.criterion = CenterLoss(len(self.classes))

        self.target_encode = TargetEncode(len(self.classes))
        self.target_decode = TargetDecode(len(self.classes))

    def work(self):
        for epoch in range(1, self.args.epochs):
            self.train(epoch)
            # self.test(epoch)

        torch.cuda.empty_cache()
        print("finish model training")

    def train(self, epoch):
        self.model.train()

        losses = []
        losses_heatmap = []
        losses_offset = []
        losses_wh = []

        pbar = tqdm(self.train_dataloader, desc=f'Test Epoch {epoch}/{self.args.epochs}',
                    bar_format='{l_bar}%s{bar}%s{r_bar}' % (Fore.BLUE, Fore.RESET))

        for data, target, line_path in pbar:
            data, target = data.to(self.device), [torch.from_numpy(ann) for ann in target]
            self.optimizer.zero_grad()
            outputs = self.model(data)
            outputs_heatmap, outputs_offset, outputs_wh = outputs[0], outputs[1], outputs[2]



            target_heatmap, target_offset, target_wh, mask,hm_mask = self.target_encode(target, outputs_heatmap, outputs_offset,
                                                                                outputs_wh)

            output = torch.cat([outputs_heatmap, outputs_offset, outputs_wh], dim=1)
            target = torch.cat([target_heatmap, target_offset, target_wh], dim=1)

            total_loss, loss_heatmap, loss_offset, loss_wh = self.criterion(output, target, mask,hm_mask)
            total_loss.backward()
            self.optimizer.step()

            losses.append(total_loss.item())
            losses_heatmap.append(loss_heatmap.item())
            losses_offset.append(loss_offset.item())
            losses_wh.append(loss_wh.item())

            pbar.set_description(
                'Train Epoch :{}/{}'.format(epoch, self.args.epochs) +
                '\t lr:{:22}'.format(self.optimizer.param_groups[0]["lr"]) +
                '\t total_loss:{:8}'.format(np.around(np.mean(losses), 6)) +
                '\t heatmap_loss:{:8}'.format(np.around(np.mean(losses_heatmap), 6)) +
                '\t offset_loss:{:8}'.format(np.around(np.mean(losses_offset), 6)) +
                '\t wh_loss:{:8}'.format(np.around(np.mean(losses_wh), 6))
            )
        self.scheduler.step()

        global loss

        if loss > np.mean(losses):
            loss = np.mean(losses)
            torch.save({"model_state_dict": self.model.state_dict()}, "weight/best.pth")
            print("model saved")

    @torch.no_grad()
    def experiment(self, epoch=1):
        self.model.eval()

        pbar = tqdm(self.train_dataloader, desc=f'Test Epoch {epoch}/{self.args.epochs}',
                    bar_format='{l_bar}%s{bar}%s{r_bar}' % (Fore.BLUE, Fore.RESET))

        for data, targets, line_path in pbar:
            data, targets = data.to(self.device), [torch.from_numpy(ann) for ann in targets]

            self.optimizer.zero_grad()
            outputs = self.model(data)
            outputs_heatmap, outputs_offset, outputs_wh = outputs[0], outputs[1], outputs[2]

            target_heatmap, target_offset, target_wh, mask = self.target_encode(targets, outputs_heatmap,
                                                                                outputs_offset,
                                                                                outputs_wh)

            target = torch.cat([target_heatmap, target_offset, target_wh], dim=1)

            outputs = self.target_decode(target)

            t_map = torch.zeros(target_heatmap.shape[2], target_heatmap.shape[3], device=target.device)
            for t in target_heatmap[0]:
                t_map += t

            t_map = t_map.cpu().numpy()
            t_image = Image.fromarray(t_map * 255)
            t_image = t_image.resize([target_heatmap.shape[3] * 4, target_heatmap.shape[2] * 4])

            image = get_image(data)
            draw = ImageDraw.Draw(image)
            if outputs[0] is None:
                continue

            for output in outputs[0]:
                x1 = output[0]
                y1 = output[1]
                x2 = output[2]
                y2 = output[3]

                draw.rectangle((x1, y1, x2, y2), width=3)
                offset = 4
                label_content = '{} {:.4f}'.format(self.classes[int(output[-2])], 1)
                label_size = draw.textsize(label_content, font)
                if y1 - label_size[1] >= 0:
                    text_origin = np.array([x1 + offset, y1 + offset - label_size[1]])
                else:
                    text_origin = np.array([x1 + offset, y1 + offset + 1])

                draw.text(text_origin, label_content, fill="black",
                          font=font)

            t_image = t_image.convert("RGBA")
            image = image.convert("RGBA")
            img = Image.blend(image, t_image, 0.3)
            img.show()
            # print("Dsa")


if __name__ == "__main__":
    model = CenterNetTrain()
    model.work()
    # model.experiment()
